package com.consume.controller;

import brave.Tracer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.ThreadLocalCurrentTraceContext;
import brave.propagation.TraceContext;
import io.micrometer.core.lang.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Configuration;

/**
 * @author zhangbowen@heytea.com
 * @since 2019-08-05 12:45
 */
@Configuration
@Slf4j
public class TraceScope implements ApplicationContextAware {
    private static ApplicationContext applicationContext;

    /**
     * 父子线程共用traceId
     *
     * @param traceContext
     * @return
     */
    public static CurrentTraceContext.Scope newThreadCrossScope(TraceContext traceContext) {
        CurrentTraceContext.Scope scope = null;
        try {
            MDC.put("traceId", traceContext.traceIdString());
            scope = CurrentTraceContext.Default.create().newScope(traceContext);
        } catch (Exception e) {
            log.warn("父子线程共用traceId，获取zipkin的scope失败");
        }
        return scope;
    }

    /**
     * 每个子线程开启新的traceId
     */
    public static CurrentTraceContext.Scope newScope(TraceContext parentTraceContext) {
        CurrentTraceContext.Scope scope = null;
        TraceContext traceContext;
        try {
            Tracer tracer = applicationContext.getBean(Tracer.class);
            if(parentTraceContext!=null) {
                traceContext = tracer.newTrace().context().toBuilder().spanId(parentTraceContext.spanId())
                        .extra(parentTraceContext.extra()).parentId(parentTraceContext.parentId()).build();
            }else {
                traceContext=tracer.newTrace().context();
            }
            MDC.put("traceId", traceContext.traceIdString());
            scope = CurrentTraceContext.Default.create().newScope(traceContext);
        } catch (Exception e) {
            log.warn("子线程开启新的traceId，获取zipkin的scope失败");
        }
        return scope;
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        TraceScope.applicationContext = applicationContext;
    }
}
